package com.ht.modules.component.ai.embedding;

import com.ht.modules.component.ai.assistants.Assistant;
import com.ht.modules.component.ai.assistants.IAssistantService;
import com.rometools.utils.Strings;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Service
@Slf4j
public class EmbeddingService {
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private IAssistantService<Assistant> assistantService;

    /**
     * 文件转换
     */
    public File convertMultiPartToFile(MultipartFile file) {
        try {
            File convFile = new File(Objects.requireNonNull(file.getOriginalFilename()));
            FileOutputStream fos = new FileOutputStream(convFile);
            fos.write(file.getBytes());
            fos.close();
            return convFile;
        } catch (Exception e) {
            throw new RuntimeException("文件转换异常:::", e);
        }
    }

    /**
     * 文件向量化存储
     */
    public void fileEmbedding(MultipartFile multipartFile) {
        log.info("==========向量化存储开始==========");
        File file = convertMultiPartToFile(multipartFile);
        String path = file.getPath();
        long l1 = System.currentTimeMillis();
        Document document = FileSystemDocumentLoader.loadDocument(path, new ApacheTikaDocumentParser());
        String text = document.text();
        String[] split = text.split("@{3}\r\n");
        List<TextSegment> segments = new ArrayList<>();
        for (int i = 0; i < split.length; i++) {
            String[] split1 = split[i].split("\r\n");
            Metadata metadata = new Metadata();
            metadata.put("content", split1[0].trim());
            metadata.put("index", i);
            StringBuilder content = new StringBuilder();
            for (int i1 = 0; i1 < split1.length; i1++) {
                if (i1 == 0) {
                    content.append("\t");
                    continue;
                }
                content.append(i1).append(".").append(split1[i1]).append("\r\n\t");
            }
            String key = split1[0].trim() + "\n" + content;
            metadata.put("summary", key);
            TextSegment textSegment = new TextSegment(key, metadata);
            segments.add(textSegment);
        }
        long l2 = System.currentTimeMillis();
        log.info("1.拆分块数:::" + segments.size() + ";耗时:::" + ((l2 - l1) / 1000) + "ms");
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        long l3 = System.currentTimeMillis();
        log.info("2.向量化数量:::" + embeddings.size() + ";耗时:::" + ((l3 - l2) / 1000) + "ms");
        EmbeddingStore<TextSegment> embeddingStore = assistantService.buildEmbeddingStore("disease");
        List<String> ids = embeddingStore.addAll(embeddings, segments);
        long l4 = System.currentTimeMillis();
        log.info("3.存储数量:::" + ids.size() + ";耗时:::" + ((l4 - l3) / 1000) + "ms");
        log.info("==========向量化存储结束==========");
        boolean exists = file.exists();
        if (exists) {
            boolean delete = file.delete();
            if (delete) {
                log.info("==========移除零时文件==========");
            } else {
                log.info("==========移除零时文件失败==========");
            }
        }
    }


    /**
     * 向量检索
     *
     * @param content 内容
     * @return 检索结果
     */
    public List<String> vectorStoreQuery(String content, String tableName,double minScore,int maxResults) {
        if (Strings.isBlank(content)) {
            throw new RuntimeException("检索内容不能为空");
        }
        if (Strings.isBlank(tableName)) {
            throw new RuntimeException("检索表名不能为空");
        }
        Embedding embedding = this.embeddingModel.embed(content).content();
        EmbeddingStore<TextSegment> embeddingStore = assistantService.buildEmbeddingStore(tableName);
        List<EmbeddingMatch<TextSegment>> embeddingMatcheList = embeddingStore.search(EmbeddingSearchRequest.builder()
                        .queryEmbedding(embedding)
                        .minScore(minScore)
                        .maxResults(maxResults)
                        .build())
                .matches();
        if (embeddingMatcheList.isEmpty()) {
            return new ArrayList<>();
        }
        return embeddingMatcheList.stream().map((item) -> {
            TextSegment textSegment = item.embedded();
            Metadata metadata = textSegment.metadata();
            Map<String, Object> map = metadata.toMap();
            return String.valueOf(map.get("summary"));
        }).toList();
    }
}
